package com.lw.leetcode.back.b;

import java.util.Arrays;

/**
 * Created with IntelliJ IDEA.
 * b
 * arr
 * https://leetcode.cn/contest/cnunionpay-2022spring/problems/I4mOGz/
 * 银联-03. 理财产品
 *
 * @author liw
 * @version 1.0
 * @date 2023/2/10 15:49
 */
public class MaxInvestment {

    public int maxInvestment(int[] product, int limit) {
        long sum = 0;
        Arrays.sort(product);
        int length = product.length;
        long max = product[length - 1];
        int count = 1;
        for (int i = length - 2; i >= 0; i--) {
            long num = product[i];
            if (num == max) {
                count++;
                continue;
            }
            long v = (max - num) * count;
            if (v <= limit) {
                limit -= v;
                sum = (sum + (((max + num + 1) * (max - num) * count) >> 1)) % 1000000007;
            } else {
                int a = limit / count;
                int b = limit % count;
                sum = (sum + (((max + max - a + 1) * a * count) >> 1) + (max - a) * b) % 1000000007;
                limit = 0;
            }
            if (limit == 0) {
                break;
            }
            count++;
            max = num;
        }
        if (limit == 0) {
            return (int) sum;
        }
        int a = limit / count;
        int b = limit % count;
        if (max - a <= 0) {
            return (int) ((sum + (((max + 1) * max * count) >> 1)) % 1000000007);
        }
        return (int) ((sum + (((max + max - a + 1) * a * count) >> 1) + (max - a) * b) % 1000000007);
    }

}
